{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Global Uncertainty Analysis: Polynomial Chaos Expansion (PCE) for Chemical Reaction Systems\n",
    "\n",
    "\n",
    "This IPython notebook uses MUQ as a basis for adaptive Polynomial Chaos Expansions to perform global uncertainty analysis for chemical reaction systems.  This IPython notebook details a workflow using RMG, Cantera, and MUQ codes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "from rmgpy.chemkin import load_chemkin_file\n",
    "from rmgpy.species import Species\n",
    "from rmgpy.tools.canteramodel import Cantera, get_rmg_species_from_user_species\n",
    "from rmgpy.tools.globaluncertainty import ReactorPCEFactory\n",
    "from rmgpy.tools.uncertainty import Uncertainty"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial setup\n",
    "\n",
    "This section sets up everything needed to perform the global uncertainty analysis. This includes creating an instance of the Uncertainty class, loading the model to be analyzed, and setting up the Cantera reactor simulator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Must use annotated chemkin file\n",
    "chemkin_file = './data/parse_source/chem_annotated.inp'\n",
    "dict_file = './data/parse_source/species_dictionary.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the Uncertainty class instance and load the model\n",
    "uncertainty = Uncertainty()\n",
    "uncertainty.load_model(chemkin_file, dict_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Map the species to the objects within the Uncertainty class\n",
    "ethane = Species(smiles='CC')\n",
    "C2H4 = Species(smiles='C=C')\n",
    "mapping = get_rmg_species_from_user_species([ethane, C2H4], uncertainty.species_list)\n",
    "\n",
    "# Define the reaction conditions\n",
    "reactor_type_list = ['IdealGasConstPressureTemperatureReactor']\n",
    "mol_frac_list = [{mapping[ethane]: 1.0}]\n",
    "Tlist = ([1300], 'K')\n",
    "Plist = ([1], 'bar')\n",
    "reaction_time_list = ([0.5], 'ms')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Global uncertainty analysis works by simulating the full model at random points within the uncertainty distributions of the input parameters. In the current implementation, the simulation is performed by Cantera, which we set up here using the RMG wrapper class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the cantera model\n",
    "job = Cantera(species_list=uncertainty.species_list, reaction_list=uncertainty.reaction_list)\n",
    "# Load the cantera model based on the RMG reactions and species\n",
    "job.load_model()\n",
    "# Generate the conditions based on the settings we declared earlier\n",
    "job.generate_conditions(reactor_type_list, reaction_time_list, mol_frac_list, Tlist, Plist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we need to load the RMG-database into the Uncertainty instance which was created in order to extract the original sources for every estimated parameter in the model. Note that the RMG-database version must be the same as the version used to generate the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uncertainty.load_database(\n",
    "    thermo_libraries=['DFT_QCI_thermo', 'primaryThermoLibrary'],\n",
    "    kinetics_families='default',\n",
    "    reaction_libraries=[],\n",
    ")\n",
    "uncertainty.extract_sources_from_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Global uncertainty analysis for uncorrelated parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assign uncorrelated parameter uncertainties \n",
    "uncertainty.assign_parameter_uncertainties(correlated=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Input a set of kinetic $(k)$ and thermo $(G)$ parameters to be propagated and their uncertainties $(\\Delta\\ln k, \\Delta G)$ into the `ReactorPCEFactory` class. These kinetic and thermo parameters should typically be pre-screened from local uncertainty analysis to narrow down to the most influential parameters.\n",
    "\n",
    "Parameter uncertainties are assigned the same way as for local uncertainty analysis and are provided directly from the `Uncertainty` instance.\n",
    "\n",
    "Random sampling from the uncertainty distributions of the input parameters is aided by a set uncertainty factors, $f$, calculated from the input uncertainties $(\\Delta\\ln k, \\Delta G)$, and a set of unit random variables, $\\xi$, sampled from a uniform distribution.\n",
    "\n",
    "For thermochemistry,\n",
    "\n",
    "$$f^G = G_{max} - G_0 = G_{0} - G_{min} = \\sqrt{3} \\Delta G$$\n",
    "\n",
    "$$G = \\xi f^G_{n} + G_{0}$$\n",
    "\n",
    "For kinetics,\n",
    "\n",
    "$$f^k = \\log_{10} \\left(\\frac{k_{max}}{k_0}\\right) = \\log_{10} \\left(\\frac{k_0}{k_{min}}\\right) = \\frac{\\sqrt{3}}{\\ln 10} \\Delta \\ln k$$\n",
    "\n",
    "$$k = 10^{\\xi f_{m}} k_{0}$$\n",
    "\n",
    "This allows calculation of a new parameter value given the nominal value, standard deviation, and the random variable.\n",
    "\n",
    "The MIT Uncertainty Quantification Library (MUQ) is used to perform the random sampling and construct a Polynomial Chaos Expansion (PCE) to fit the output variable of interest, mole fractions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose input parameters to vary within their uncertainty bounds\n",
    "k_params = [3, 5]  # RMG indices of reactions to vary\n",
    "g_params = [1, 4]  # RMG indices of species to vary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create ReactorPCEFactory global uncertainty analysis object for the uncorrelated case\n",
    "reactor_pce_factory = ReactorPCEFactory(\n",
    "    cantera=job,\n",
    "    output_species_list=[mapping[ethane], mapping[C2H4]],\n",
    "    k_params=k_params,\n",
    "    k_uncertainty=uncertainty.kinetic_input_uncertainties,\n",
    "    g_params=g_params,\n",
    "    g_uncertainty=uncertainty.thermo_input_uncertainties,\n",
    "    correlated=False,\n",
    "    logx=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Begin generating the PCEs adaptively based a runtime. If you would like to re-generate the PCEs with different settings, the previous block must also be re-run.\n",
    "\n",
    "There are actually four methods for generating PCEs. See the `ReactorPCEFactory.generate_pce` function for more details.\n",
    "\n",
    "- Option 1: Adapt PCE for a specified amount of time\n",
    "- Option 2: Adapt PCE until error tolerance is met\n",
    "- Option 3: Adapt PCE until specified number of model evaluations\n",
    "- Option 4: Used a fixed order, and (optionally) adapt later.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reactor_pce_factory.generate_pce(\n",
    "    max_adapt_time=60,  # Adapt for 1 minute; note that this will not be very accurate\n",
    "    error_tol=None,\n",
    "    max_evals=None,\n",
    "    should_adapt=True,  # If set to false, will use a fixed order\n",
    "    start_order=2,\n",
    ")\n",
    "\n",
    "print('Number of Model Evaluations: {0}'.format(reactor_pce_factory.factory.NumEvals()))\n",
    "print('Estimated L2 Error: {0:.4e}'.format(reactor_pce_factory.factory.Error()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare the outputs for a test point using the real model versus using the PCE approximation.\n",
    "Evaluate the desired output mole fractions based on a set of inputs `inputs = [[ln(k)_rv], [g_rv]]` which contains the \n",
    "random unit uniform variables attributed to the uncertain kinetics and free energy parameters, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a random test point of length = number of k_params + number of g_params\n",
    "random_test_point = [random.uniform(-1.0,1.0) for i in range(len(k_params)+len(g_params))]\n",
    "true_test_point_output, pce_test_point_output = reactor_pce_factory.compare_output(random_test_point, log=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Obtain the results: the species mole fraction mean and variance computed from the PCE, as well as the global sensitivity indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean, variance, covariance, main_sens, total_sens = reactor_pce_factory.analyze_results(log=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Global uncertainty analysis of correlated parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uncertainty.assign_parameter_uncertainties(correlated=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k_params = [\n",
    "    'Training H_Abstraction H(6)+ethane(1)=H2(11)+C2H5(5)',\n",
    "    'Training R_Recombination CH3(4)+CH3(4)=ethane(1)',\n",
    "]\n",
    "g_params = [\n",
    "    'Library CH4(3) ',\n",
    "    'Estimation CH3(4)',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reactor_pce_factory_correlated = ReactorPCEFactory(\n",
    "    cantera=job,\n",
    "    output_species_list=[mapping[ethane], mapping[C2H4]],\n",
    "    k_params=k_params,\n",
    "    k_uncertainty=uncertainty.kinetic_input_uncertainties,\n",
    "    g_params=g_params,\n",
    "    g_uncertainty=uncertainty.thermo_input_uncertainties,\n",
    "    correlated=True,\n",
    "    logx=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do the same analysis for the correlated `reactorPCEFactory`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reactor_pce_factory_correlated.generate_pce(\n",
    "    max_adapt_time=60,  # Adapt for 1 minute; note that this will not be very accurate\n",
    "    error_tol=None,\n",
    "    max_evals=None,\n",
    "    should_adapt=True,  # If set to false, will use a fixed order\n",
    "    start_order=2,\n",
    ")\n",
    "\n",
    "print('Number of Model Evaluations: {0}'.format(reactor_pce_factory_correlated.factory.NumEvals()))\n",
    "print('Estimated L2 Error: {0:.4e}'.format(reactor_pce_factory_correlated.factory.Error()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_test_point = [random.uniform(-1.0,1.0) for i in range(len(k_params)+len(g_params))]\n",
    "true_test_point_output, pce_test_point_output = reactor_pce_factory_correlated.compare_output(random_test_point, log=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean, variance, covariance, main_sens, total_sens = reactor_pce_factory_correlated.analyze_results(log=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:rmg_env]",
   "language": "python",
   "name": "conda-env-rmg_env-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
